#include <iostream>
#include <cassert>
#include <vector>
#include <thread>
#include <random>
#include <algorithm>

#include "../code/ThreeLayersMemoryPool.hpp"

// 基础分配测试
void testBasicAllocation()
{
    std::cout << "Running basic allocation test..." << std::endl;

    // 测试小内存分配
    void *ptr1 = MemoryPool::ThreeLayersMemoryPool::alloc(8);
    assert(ptr1 != nullptr);
    MemoryPool::ThreeLayersMemoryPool::dealloc(ptr1, 8);

    // 测试中内存分配
    void *ptr2 = MemoryPool::ThreeLayersMemoryPool::alloc(1024);
    assert(ptr2 != nullptr);
    MemoryPool::ThreeLayersMemoryPool::dealloc(ptr2, 1024);

    // 测试大内存分配（超过Common.h SizeManager::MAX_ALLOC_BYTES 256 * 1024）
    void *ptr3 = MemoryPool::ThreeLayersMemoryPool::alloc(1024 * 1024);
    assert(ptr3 != nullptr);
    MemoryPool::ThreeLayersMemoryPool::dealloc(ptr3, 1024 * 1024);

    std::cout << "Basic allocation test passed!" << std::endl;
}

// 内存写入测试
void testMemoryWriting()
{
    std::cout << "Running memory writing test..." << std::endl;

    // 分配内存
    const size_t NEED_SIZE = 128;
    char *ptr = static_cast<char *>(MemoryPool::ThreeLayersMemoryPool::alloc(NEED_SIZE));
    assert(ptr != nullptr);

    // 写入数据
    for (size_t i = 0; i < NEED_SIZE; ++i)
    {
        ptr[i] = static_cast<char>(i % 256);
    }

    // 验证数据
    for (size_t i = 0; i < NEED_SIZE; ++i)
    {
        assert(ptr[i] == static_cast<char>(i % 256));
    }

    MemoryPool::ThreeLayersMemoryPool::dealloc(ptr, NEED_SIZE);

    std::cout << "Memory writing test passed!" << std::endl;
}

// 多线程测试
void testMultiThreading()
{
    std::cout << "Running multi-threading test..." << std::endl;

    const int THREAD_ALLOC_NUM = 1000; // 每线程分配次数
    std::atomic<bool> has_error{false};
    const int THREAD_NUM = 4;

    // 常量不捕获
    auto threadFunc = [&has_error]()
    {
        try
        {
            std::vector<std::pair<void *, size_t>> allocs; // 记录分配的指针-大小，用于释放
            allocs.reserve(THREAD_ALLOC_NUM);

            // 分配
            for (int i = 0; i < THREAD_ALLOC_NUM && !has_error; ++i)
            {
                size_t size = (rand() % 256 + 1) * 8;
                void *ptr = MemoryPool::ThreeLayersMemoryPool::alloc(size);
                if (!ptr)
                {
                    std::cerr << "Allocation failed for size: " << size << std::endl;
                    has_error = true;
                    break;
                }

                allocs.push_back({ptr, size});

                // 随机释放
                if (rand() % 2 && !allocs.empty())
                {
                    size_t index = rand() % allocs.size();
                    MemoryPool::ThreeLayersMemoryPool::dealloc(allocs[index].first, allocs[index].second);
                    allocs.erase(allocs.begin() + index);
                }
            }

            // 全部释放
            for (const auto &alloc : allocs)
            {
                MemoryPool::ThreeLayersMemoryPool::dealloc(alloc.first, alloc.second);
            }
        }
        catch (const std::exception &e)
        {
            std::cerr << "Thread exception: " << e.what() << std::endl;
            has_error = true;
        }
    };

    // 启动线程
    std::vector<std::thread> threads{};
    for (int i = 0; i < THREAD_NUM; ++i)
    {
        threads.emplace_back(threadFunc);
    }

    for (auto &thread : threads)
    {
        thread.join();
    }

    std::cout << "Multi-threading test passed!" << std::endl;
}

// 边界测试
void testEdgeCase()
{
    std::cout << "Running edge cases test..." << std::endl;

    // 0
    void *ptr1 = MemoryPool::ThreeLayersMemoryPool::alloc(0);
    assert(ptr1 != nullptr);
    MemoryPool::ThreeLayersMemoryPool::dealloc(ptr1, 0);

    // 1->最小对齐
    void *ptr2 = MemoryPool::ThreeLayersMemoryPool::alloc(1);
    assert(ptr2 != nullptr);
    // 检查地址是否对齐
    // 内存地址较大整数，uintptr_t无符号整数类型，足以存储指针的值并进行位操作，假设prt2八字节对齐存储0x1008
    // 8 - 1 = 7 = 0b111
    // 0000 1000 & 0000 0111 == 0000，证明是8B对齐的，其他同理
    // assert()返回非0继续，0终止程序+错误信息
    assert((reinterpret_cast<uintptr_t>(ptr2) & (MemoryPool::SizeManager::ALIGNMENT_BYTES - 1)) == 0);
    MemoryPool::ThreeLayersMemoryPool::dealloc(ptr2, 1);

    // 最大
    void *ptr3 = MemoryPool::ThreeLayersMemoryPool::alloc(MemoryPool::SizeManager::MAX_ALLOC_BYTES);
    assert(ptr3 != nullptr);
    MemoryPool::ThreeLayersMemoryPool::dealloc(ptr3, MemoryPool::SizeManager::MAX_ALLOC_BYTES);

    // 超过最大
    void *ptr4 = MemoryPool::ThreeLayersMemoryPool::alloc(MemoryPool::SizeManager::MAX_ALLOC_BYTES + 1);
    assert(ptr4 != nullptr);
    MemoryPool::ThreeLayersMemoryPool::dealloc(ptr4, MemoryPool::SizeManager::MAX_ALLOC_BYTES + 1);

    std::cout << "Edge cases test passed!" << std::endl;
}

// 压力测试
void testStress()
{
    std::cout << "Running stress test..." << std::endl;

    const int ITERATION_NUM = 10000;
    std::vector<std::pair<void *, size_t>> allocs;
    allocs.reserve(ITERATION_NUM);

    for (int i = 0; i < ITERATION_NUM; ++i)
    {
        size_t size = (rand() % 1024 + 1) * 8;
        void *ptr = MemoryPool::ThreeLayersMemoryPool::alloc(size);
        assert(ptr != nullptr);
        allocs.push_back({ptr, size});
    }

    // 随机顺序释放
    std::random_device rd{};
    std::mt19937 gen(rd());
    std::shuffle(allocs.begin(), allocs.end(), gen);
    for (const auto &alloc : allocs)
    {
        MemoryPool::ThreeLayersMemoryPool::dealloc(alloc.first, alloc.second);
    }

    std::cout << "Stress test passed!" << std::endl;
}

int main()
{
    try
    {
        std::cout << "Starting memory pool tests..." << std::endl;

        testBasicAllocation();
        testMemoryWriting();
        testMultiThreading();
        testEdgeCase();
        testStress();

        std::cout << "All tests passed successfully!" << std::endl;
        return 0;
    }
    catch (const std::exception &e)
    {
        std::cerr << "Test failed with exception: " << e.what() << std::endl;
        return 1;
    }
}